import psutil
import socket
import os
from Tool.SelectData import CzIp

def get_process_cmdline(pid):
    try:
        with open(f"/proc/{pid}/cmdline", "rb") as f:
            cmdline = f.read().decode().replace('\x00', ' ')
            exe_index = cmdline.find('.exe')
            if exe_index != -1:
                program_path = cmdline[:exe_index + 4]
                return program_path.strip()
    except FileNotFoundError:
        return None

def get_network_connections():
    netstat = {}
    connections = psutil.net_connections()
    for conn in connections:
        本地地址 = f"{conn.laddr.ip}:{conn.laddr.port}"

        远程地址 = "0.0.0.0"
        if conn.raddr:
            远程地址 = f"{conn.raddr.ip}:{conn.raddr.port}"

        if conn.type == socket.SOCK_STREAM:
            协议 = "TCP"
        elif conn.type == socket.SOCK_DGRAM:
            协议 = "UDP"
        else:
            协议 = "UNKNOWN"

        if conn.family == socket.AF_INET:
            地址族 = "IPv4"
        elif conn.family == socket.AF_INET6:
            地址族 = "IPv6"
        else:
            地址族 = "UNKNOWN"

        连接状态 = conn.status

        进程号 = conn.pid
        进程名称 = "-"
        程序路径 = "-"
        if 进程号 is not None:
            try:
                进程名称 = psutil.Process(进程号).name()
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                进程名称 = "-"
            try:
                程序路径 = os.readlink(f"/proc/{进程号}/exe")
                if 'wine' in 程序路径:
                    Win路径 = get_process_cmdline(进程号)
                    if 'wine' not in Win路径:
                        程序路径 = Win路径.capitalize()
            except:
                  pass
            
        
        try:
            cz = CzIp()
            addres = cz.get_addr_by_ip(远程地址.split(':')[0])
            远程地址归属地 = f'{addres[0]} {addres[1]}'
        except Exception as e:
            远程地址归属地 = '-'

        字典名称 = 进程名称.split(' ')[0]
        if 字典名称 in netstat:
            netstat[字典名称].append([进程名称,进程号,程序路径,协议,本地地址,远程地址,连接状态,远程地址归属地])
        else:
            netstat[字典名称] = []
            netstat[字典名称].append([进程名称,进程号,程序路径,协议,本地地址,远程地址,连接状态,远程地址归属地])
    return netstat
